/*
 * Functions to interract with virtiofs-pci
 */
#include "precomp.h"
#pragma hdrstop

#define INITGUID
#include "mrxglobs.h"
#include "viofs_call.h"
#include "../../viofs/shared/fuse.h"
#include "fuse_structs.h"

// Allow use of ExAllocatePoolWithTag()
#pragma warning(disable:4996)

NTSTATUS
OpenViofsDevice(
    IN  PHWFS_DEVICE_EXTENSION  pDevExt)
{
    PWSTR symbolicLink;
    PWSTR symbolicLinkList = NULL;
    PFILE_OBJECT fileObject;
    PDEVICE_OBJECT devObject;
    NTSTATUS Status = STATUS_SUCCESS;

    PIRP Irp;
    KEVENT Event;
    IO_STATUS_BLOCK IoStatus = { 0 };
    VIRTIOFS_INKERNEL_FUSE_REQUEST_ROUTINE ViofsRequestParam;

    HWFS_DBG_PRINT(DBG_TRACE, DBG_SESSION, ("OpenViofsDevice\n"));

    if (pDevExt->ViofsFileObject) {
        HWFS_DBG_PRINT(DBG_TRACE, DBG_SESSION, ("Device already opened\n"));
        goto finish;
    }

    Status = IoGetDeviceInterfaces(
        &GUID_DEVINTERFACE_VIRT_FS,
        NULL,
        0,
        &symbolicLinkList);
    if (!NT_SUCCESS(Status) || NULL == symbolicLinkList) {
        HWFS_DBG_PRINT(DBG_ERR, DBG_ANY,
            ("Failed to enumerate virtiofs: 0x%x\n", Status));
        if (NT_SUCCESS(Status)) {
            Status = STATUS_UNSUCCESSFUL;
        }
        goto finish;
    }

    for (symbolicLink = symbolicLinkList;
            symbolicLink[0] != L'\0' && symbolicLink[1] != L'\0';
            symbolicLink += wcslen(symbolicLink) + 1) {
        UNICODE_STRING objName;
        RtlInitUnicodeString(&objName, symbolicLink);

        Status = IoGetDeviceObjectPointer(
            &objName,
            FILE_ALL_ACCESS,
            &fileObject,
            &devObject
        );
        if (NT_SUCCESS(Status)) {
            HWFS_DBG_PRINT(DBG_TRACE, DBG_SESSION, ("Successfully opened\n"));
            pDevExt->ViofsFileObject = fileObject;
            pDevExt->ViofsDeviceObject = devObject;
            break;
        } else {
            HWFS_DBG_PRINT(DBG_ERR, DBG_ANY,
                ("Failed to open virtiofs: 0x%x\n", Status));
            continue;
        }
    }

    if (NULL == pDevExt->ViofsFileObject) {
        Status = STATUS_UNSUCCESSFUL;
        goto finish;
    }

    // Obtain pointer to ViofsSubmitRequest

    Status = STATUS_SUCCESS;

    KeInitializeEvent(&Event, NotificationEvent, FALSE);
    ViofsRequestParam.ApiVersion = 0xffffffff;
    Irp = IoBuildDeviceIoControlRequest(
        IOCTL_VIRTFS_GET_INKERNEL_API,
        pDevExt->ViofsDeviceObject,
        NULL, 0,
        &ViofsRequestParam, sizeof(ViofsRequestParam),
        FALSE,
        &Event,
        &IoStatus
    );
    if (!Irp) {
        HWFS_DBG_PRINT(DBG_ERR, DBG_ANY, ("Failed to alloc Irp for virtiofs\n"));
        Status = STATUS_INSUFFICIENT_RESOURCES;
    } else {
        if (IoCallDriver(pDevExt->ViofsDeviceObject, Irp) == STATUS_PENDING) {
            KeWaitForSingleObject(&Event, Executive, KernelMode, FALSE, NULL);
        }

        if (IoStatus.Status != 0) {
            Status = IoStatus.Status;
            HWFS_DBG_PRINT(DBG_ERR, DBG_ANY, ("Get Viofs Inkernel API failed: 0x%x\n", Status));
        } else if (ViofsRequestParam.ApiVersion != VIRTFS_INKERNEL_API_VERSION) {
            HWFS_DBG_PRINT(DBG_ERR, DBG_ANY,
                ("Get Viofs Inkernel API returned incorrect version: 0x%I64x\n",
                ViofsRequestParam.ApiVersion));
            Status = STATUS_UNSUCCESSFUL;
        } else {
            pDevExt->ViofsSubmitRequest = ViofsRequestParam.Func;
            pDevExt->ViofsCancelRequest = ViofsRequestParam.CancelFunc;
            pDevExt->ViofsContext = ViofsRequestParam.ViofsContext;
            pDevExt->ViofsResetDevice = ViofsRequestParam.ResetFunc;
        }
    }

    if (!NT_SUCCESS(Status)) {
        ObDereferenceObject(pDevExt->ViofsFileObject);
        pDevExt->ViofsFileObject = NULL;
        pDevExt->ViofsDeviceObject = NULL;
        goto finish;
    }

finish:
    if (symbolicLinkList != NULL) {
        ExFreePool(symbolicLinkList);
    }
    HWFS_DBG_PRINT(DBG_TRACE, DBG_SESSION, ("OpenViofsDevice: Status 0x%x\n", Status));
    return Status;
}

VOID
CloseViofsDevice(
    IN  PHWFS_DEVICE_EXTENSION  pDevExt)
{
    if (!pDevExt->ViofsFileObject) {
        return;
    }

    ObDereferenceObject(pDevExt->ViofsFileObject);
    pDevExt->ViofsFileObject = NULL;
    pDevExt->ViofsDeviceObject = NULL;
}

static
PMDL
HwFsAllocatePages(ULONG bytes)
{
    PHYSICAL_ADDRESS Low;
    PHYSICAL_ADDRESS High;
    PHYSICAL_ADDRESS Skip;
    Low.QuadPart = 0;
    High.QuadPart = -1;
    Skip.QuadPart = 0;
    return MmAllocatePagesForMdlEx(Low, High, Skip,
        bytes, MmCached,
        MM_DONT_ZERO_ALLOCATION | MM_ALLOCATE_FULLY_REQUIRED);
}

#define ROUND_TO8(Bytes) (((Bytes) + 7) & ~7ul)

static
VOID
ViofsRequestCompleteCallback(
    IN  PVIOFS_INKERNEL_REQUEST pRequest,
    IN  UINT64                  BytesRetured)
{
    KEVENT* pEvent;

    PVIOFSCALL_EXTENSION    pRequestExt = pRequest->CallerContext;

    HWFS_DBG_PRINT(DBG_TRACE, DBG_VIOFS,
        ("ViofsRequestCompleteCallback for %p; BytesReturned 0x%I64x\n",
            pRequest, BytesRetured));

    pRequestExt->BytesReturned = BytesRetured;
    KeSetEvent(&pRequestExt->CompleteEvent, 0, FALSE);
}

VOID
FreeViofsRequest(
    IN  PVIOFS_INKERNEL_REQUEST pRequest)
{
    PVIOFSCALL_EXTENSION    pRequestExt;

    pRequestExt = pRequest->CallerContext;

    ASSERT(pRequestExt != NULL);

    // both in and out are in the same buffer
    if (pRequestExt->FuseInHdr) {
        MmUnmapLockedPages(pRequestExt->FuseInHdr, pRequestExt->HdrMdl);
    }

    if (pRequestExt->HdrMdl) {
        MmFreePagesFromMdl(pRequestExt->HdrMdl);
        ExFreePool(pRequestExt->HdrMdl);
    }

    if (pRequestExt->MdlsAllocated) {
        if (pRequestExt->InputBufferMdl) {
            IoFreeMdl(pRequestExt->InputBufferMdl);
        }
        if (pRequestExt->OutputBufferMdl) {
            IoFreeMdl(pRequestExt->OutputBufferMdl);
        }
    }

    ExFreePoolWithTag(pRequest, 'QRFV');
}

// Allocates request capable to hold fuse_in/out header
// and buffers with sg_count
static
PVIOFS_INKERNEL_REQUEST
AllocViofsRequestInternal(
    IN  PHWFS_DEVICE_EXTENSION  pDevExt,
    IN  ULONG  FuseInHdrSize,
    IN  ULONG  FuseOutHdrSize,
    IN  ULONG  sg_count)
{
    PUCHAR   cp;
    PUCHAR   pHdrBuf;
    ULONG    HdrBufSize;
    ULONG    TotalSgs;
    PVIOFS_INKERNEL_REQUEST pRequest;
    PVIOFSCALL_EXTENSION    pRequestExt;

    ULONG CtxSize = sizeof(VIOFSCALL_EXTENSION);

    TotalSgs = BYTES_TO_PAGES(FuseInHdrSize) + BYTES_TO_PAGES(FuseOutHdrSize) + sg_count;

    pRequest = ExAllocatePoolWithTag(
        NonPagedPool,
        sizeof(VIOFS_INKERNEL_REQUEST) + TotalSgs * sizeof(VIOFS_SG) + ROUND_TO8(CtxSize),
        'QRFV');
    if (NULL == pRequest) {
        return NULL;
    }

    RtlZeroMemory(pRequest, sizeof(*pRequest));

    pRequestExt = (PVIOFSCALL_EXTENSION)(pRequest + 1);
    pRequest->CallerContext = pRequestExt;
    RtlZeroMemory(pRequestExt, sizeof(*pRequestExt));

    //
    // Don't bother and allocate 2 pages: one for in-header, other for out-header.
    //
    HdrBufSize = ROUND_TO_PAGES(FuseInHdrSize) + ROUND_TO_PAGES(FuseOutHdrSize);
    pRequestExt->HdrMdl = HwFsAllocatePages(HdrBufSize);
    if (NULL == pRequestExt->HdrMdl) {
        FreeViofsRequest(pRequest);
        return NULL;
    }

    pHdrBuf = MmMapLockedPagesSpecifyCache(pRequestExt->HdrMdl, KernelMode,
        MmCached, NULL, FALSE, NormalPagePriority | MdlMappingNoExecute);
    if (NULL == pHdrBuf) {
        FreeViofsRequest(pRequest);
        return NULL;
    }
    pRequestExt->FuseInHdr = (FUSE_IN_HEADER *)pHdrBuf;
    pRequestExt->FuseOutHdr = (FUSE_OUT_HEADER *)(pHdrBuf + ROUND_TO_PAGES(FuseOutHdrSize));
    pRequestExt->InHdrSize = FuseInHdrSize;
    pRequestExt->OutHdrSize = FuseOutHdrSize;

    pRequest->Sg = (PVIOFS_SG)(((PUCHAR)pRequestExt) + ROUND_TO8(CtxSize));

    pRequest->CompleteCallback = ViofsRequestCompleteCallback;

    return pRequest;
}


static ULONG TotalSgs(ULONG Bytes)
{
    return 1 + BYTES_TO_PAGES(Bytes);
}

static
ULONG
FillSgFromMdl(
    OUT PVIOFS_SG  sg,
    IN  PMDL       Mdl,
    IN  ULONG      Length,
    IN  ULONG      SkipPages)
{
    ULONG SgCount;
    ULONG SkipCount;
    ULONG PageOffset;

    if (Length == 0) {
        return 0;
    }

    SgCount = 0;
    SkipCount = 0;
    // first buffer in Mdl can be non-paged aligned
    PageOffset = MmGetMdlByteOffset(Mdl);
    while (Mdl != NULL) {
        ULONG i;
        PPFN_NUMBER pfn_array;
        ULONG MdlByteCount = MmGetMdlByteCount(Mdl);
        ULONG total_pages = BYTES_TO_PAGES(PageOffset + MdlByteCount);

        pfn_array = MmGetMdlPfnArray(Mdl);
        for (i = 0; i < total_pages; i++) {
            if (SkipCount < SkipPages) {
                SkipCount++;
                PageOffset = 0;
                continue;
            }

            ULONG ChunkSize = PAGE_SIZE - PageOffset;
            ULONG len = (ULONG)(min(Length, ChunkSize));
            Length -= len;
            sg[SgCount].PhyAddr.QuadPart = ((ULONGLONG)pfn_array[i] << PAGE_SHIFT) + PageOffset;
            PageOffset = 0;
            sg[SgCount].Length = len;
            SgCount++;
            if (Length == 0) {
                goto finish;
            }
        }
        Mdl = Mdl->Next;
    }
finish:
    return SgCount;
}

static
VOID
ViofsRequestFillSgs(
    IN PVIOFS_INKERNEL_REQUEST pRequest,
    IN ULONG                   MdlSkipPages)
{
    PVIOFS_SG pCurrSg;
    ULONG Count;
    PVIOFSCALL_EXTENSION    pRequestExt;

    pRequestExt = pRequest->CallerContext;

    pCurrSg = pRequest->Sg;

    //
    // fuse_in_header
    //
    Count = FillSgFromMdl(pCurrSg, pRequestExt->HdrMdl, pRequestExt->InHdrSize, 0);
    pRequest->InputSgCount = Count;
    pCurrSg += Count;

    //
    // InputBuffer
    //
    if (pRequestExt->InputBufferSize) {
        Count = FillSgFromMdl(
            pCurrSg, pRequestExt->InputBufferMdl, pRequestExt->InputBufferSize, MdlSkipPages);
        pCurrSg += Count;
        pRequest->InputSgCount += Count;
    }

    //
    // fuse_out_hdr
    //
    Count = FillSgFromMdl(
        pCurrSg,
        pRequestExt->HdrMdl,
        pRequestExt->OutHdrSize,
        BYTES_TO_PAGES(pRequestExt->InHdrSize) /* skip pages of FuseInHdr*/);
    pCurrSg += Count;
    pRequest->OutputSgCount = Count;

    //
    // OutputBuffer
    //
    if (pRequestExt->OutputBufferSize) {
        Count = FillSgFromMdl(
            pCurrSg,
            pRequestExt->OutputBufferMdl, pRequestExt->OutputBufferSize, MdlSkipPages);

        pCurrSg += Count;
        pRequest->OutputSgCount += Count;
    }
}

// Buffers must be allocated in kernel, in non-paged pool
PVIOFS_INKERNEL_REQUEST
AllocViofsRequestWithBuffer(
    IN  PHWFS_DEVICE_EXTENSION  pDevExt,
    IN  ULONG FuseInHdrSize,
    IN  ULONG FuseOutHdrSize,
    IN  PVOID InputBuffer,
    IN  ULONG InputBufferSize,
    IN  PVOID OutputBuffer,
    IN  ULONG OutputBufferSize)
{
    PVIOFS_INKERNEL_REQUEST pRequest;
    PVIOFSCALL_EXTENSION    pRequestExt;

    pRequest = AllocViofsRequestInternal(
        pDevExt,
        FuseInHdrSize, FuseOutHdrSize,
        TotalSgs(InputBufferSize) + TotalSgs(OutputBufferSize));
    if (NULL == pRequest) {
        return NULL;
    }

    pRequestExt = pRequest->CallerContext;

    pRequestExt->InputBufferSize = InputBufferSize;
    pRequestExt->OutputBufferSize = OutputBufferSize;

    pRequestExt->MdlsAllocated = TRUE;

    if (InputBufferSize) {
        pRequestExt->InputBufferMdl = IoAllocateMdl(InputBuffer, InputBufferSize, FALSE, FALSE, NULL);
        if (NULL == pRequestExt->InputBufferMdl) {
            goto error;
        }
        MmBuildMdlForNonPagedPool(pRequestExt->InputBufferMdl);
    }

    if (OutputBufferSize) {
        pRequestExt->OutputBufferMdl = IoAllocateMdl(OutputBuffer, OutputBufferSize, FALSE, FALSE, NULL);
        if (NULL == pRequestExt->OutputBufferMdl) {
            goto error;
        }
        MmBuildMdlForNonPagedPool(pRequestExt->OutputBufferMdl);
    }

    ViofsRequestFillSgs(pRequest, 0);

    return pRequest;

error:
    FreeViofsRequest(pRequest);
    return NULL;
}

PVIOFS_INKERNEL_REQUEST
AllocViofsRequestWithMdl(
    IN  PHWFS_DEVICE_EXTENSION  pDevExt,
    IN  ULONG FuseInHdrSize,
    IN  ULONG FuseOutHdrSize,
    IN  ULONG MdlSkipPages,
    IN  PMDL  InputBuffer,
    IN  ULONG InputBufferSize,
    IN  PMDL  OutputBuffer,
    IN  ULONG OutputBufferSize)
{
    PVIOFS_INKERNEL_REQUEST pRequest;
    PVIOFSCALL_EXTENSION    pRequestExt;

    pRequest = AllocViofsRequestInternal(
        pDevExt,
        FuseInHdrSize, FuseOutHdrSize,
        TotalSgs(InputBufferSize) + TotalSgs(OutputBufferSize));
    if (NULL == pRequest) {
        return NULL;
    }

    pRequestExt = pRequest->CallerContext;

    pRequestExt->InputBufferSize = InputBufferSize;
    pRequestExt->OutputBufferSize = OutputBufferSize;

    pRequestExt->MdlsAllocated = FALSE;
    pRequestExt->InputBufferMdl = InputBuffer;
    pRequestExt->OutputBufferMdl = OutputBuffer;

    ViofsRequestFillSgs(pRequest, MdlSkipPages);

    return pRequest;
}

static
VOID
CancelViofsRequest(
    IN  PHWFS_DEVICE_EXTENSION  pDevExt,
    IN  PVIOFS_INKERNEL_REQUEST pRequest)
{
    PVIOFSCALL_EXTENSION pRequestExt;
    BOOLEAN cancelled;

    HWFS_DBG_PRINT(DBG_WARN, DBG_VIOFS, ("Cancel Viofs Request %p\n", pRequest));
    cancelled = pDevExt->ViofsCancelRequest(pDevExt->ViofsContext, pRequest);
    if (cancelled) {
        HWFS_DBG_PRINT(DBG_WARN, DBG_VIOFS, ("Viofs Request cancelled: %p\n", pRequest));
    } else {
        HWFS_DBG_PRINT(DBG_WARN, DBG_VIOFS,
            ("Viofs Request appears to be already completed. Wait callback..\n"));

        pRequestExt = (PVIOFSCALL_EXTENSION)pRequest->CallerContext;
        KeWaitForSingleObject(&pRequestExt->CompleteEvent,
            Executive, KernelMode, FALSE, NULL);
        HWFS_DBG_PRINT(DBG_WARN, DBG_VIOFS, ("Viofs Request cancelled: %p\n", pRequest));
    }
    pDevExt->ViofsResetDevice(pDevExt->ViofsContext);
}

NTSTATUS
SubmitViofsRequest(
    IN  PHWFS_DEVICE_EXTENSION  pDevExt,
    IN  PVIOFS_INKERNEL_REQUEST pRequest)
{
    NTSTATUS Status;
    PVIOFSCALL_EXTENSION pRequestExt;
    KEVENT* pEvent;

    HWFS_DBG_PRINT(DBG_TRACE, DBG_VIOFS, ("Call Viofs with Request %p\n", pRequest));

    pRequestExt = (PVIOFSCALL_EXTENSION)pRequest->CallerContext;
    KeInitializeEvent(&pRequestExt->CompleteEvent, NotificationEvent, FALSE);

    Status = pDevExt->ViofsSubmitRequest(pDevExt->ViofsContext, pRequest);

    HWFS_DBG_PRINT(DBG_TRACE, DBG_VIOFS, ("ViofsSubmitRequest Status 0x%x\n", Status));
    if (Status == STATUS_PENDING) {
        Status = KeWaitForSingleObject(&pRequestExt->CompleteEvent,
            Executive, KernelMode, FALSE, pDevExt->FuseTimeoutPtr);
        HWFS_DBG_PRINT(DBG_TRACE, DBG_VIOFS,
            ("ViofsSubmitRequest Finished (%x); BytesReturned %u\n", Status, pRequestExt->BytesReturned));
        if (Status == STATUS_TIMEOUT) {
            // Reduce timeout to small value to minimize next wait
            pDevExt->FuseTimeout.QuadPart = -1 * SUBMIT2_VIOFS_TIMEOUT;
            CancelViofsRequest(pDevExt, pRequest);
            return STATUS_NETWORK_UNREACHABLE;
        }
        Status = STATUS_SUCCESS;
    }
    // Reset timeout to original value
    pDevExt->FuseTimeout.QuadPart = -1 * SUBMIT_VIOFS_TIMEOUT;

    if (!NT_SUCCESS(Status)) {
        HWFS_DBG_PRINT(DBG_ERR, DBG_VIOFS,
            ("ViofsSubmitError: 0x%x\n", Status));
        // Always return NETWORK_UNREACHABLE to indicate backend error
        return STATUS_NETWORK_UNREACHABLE;
    }

    if (pRequestExt->BytesReturned < sizeof(FUSE_OUT_HEADER) &&
            pRequest->FuseOpcode != FUSE_FORGET) {
        HWFS_DBG_PRINT(DBG_ERR, DBG_VIOFS,
            ("BytesReturned less then fuse_out_header: %u\n", pRequestExt->BytesReturned));
        return STATUS_DATA_OVERRUN;
    }

    return Status;
}
